from llama_index.core.agent.workflow import FunctionAgent
from llama_index.llms.openai import OpenAI
from llama_index.tools.yahoo_finance import YahooFinanceToolSpec

#使用现有工具

def multiply(a: float, b: float) -> float:
    """Multiply two numbers and returns the product"""
    return a * b


def add(a: float, b: float) -> float:
    """Add two numbers and returns the sum"""
    return a + b


finance_tools = YahooFinanceToolSpec().to_tool_list()

finance_tools.extend([multiply, add])

llm = OpenAI(model="gpt-4.1-nano",
             api_key="sk-CftUbVSsA61lwwgMz9xvt6znTunQZfgBP8ZCVLbQsKfXUR6k",
             api_base="https://www.henapi.top/v1", )

workflow = FunctionAgent(
    name="Agent",
    description="Useful for performing financial operations.",
    llm=llm,
    tools=finance_tools,
    system_prompt="You are a helpful assistant.",
)


async def main():
    response = await workflow.run(
        user_msg="What's the current stock price of NVIDIA?"
    )
    print(response)


if __name__ == "__main__":
    import asyncio

    asyncio.run(main())

# 跑了很长时间，但是没有结果出来，放弃